/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 * All rights reserved.
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */
#define FBGEMM_EXPORTS

#include "fbgemm/FbgemmEmbedding.h"

#include <asmjit/asmjit.h>
#include <cpuinfo.h>
#include <immintrin.h>
#include <cassert>
#include <cmath>
#include <iostream>
#include <map>
#include <mutex>
#include <string>
#include <tuple>
#include "./CodeCache.h"
#include "./MaskAvx2.h"
#include "./RefImplementations.h"
#include "fbgemm/Types.h"

namespace fbgemm {

namespace {

namespace x86 = asmjit::x86;

template <
    typename inType,
    typename indxType,
    typename offsetType,
    bool ROWWISE_SPARSE>
class ReturnFunctionSignature {};

template <typename inType, typename indxType, typename offsetType>
class ReturnFunctionSignature<inType, indxType, offsetType, false> {
 public:
  using jit_embedding_kernel = bool (*)(
      std::int64_t output_size,
      std::int64_t index_size,
      std::int64_t data_size,
      const inType* input,
      const indxType* indices,
      const offsetType* offsets_or_lengths,
      const float* weights,
      float* out,
      const int* mask);
};

template <typename inType, typename indxType, typename offsetType>
class ReturnFunctionSignature<inType, indxType, offsetType, true> {
 public:
  using jit_embedding_kernel = bool (*)(
      std::int64_t output_size,
      std::int64_t index_size,
      int64_t uncompressed_data_size,
      // int64_t compressed_data_size,
      const inType* input,
      const indxType* indices,
      const offsetType* offsets_or_lengths,
      const float* weights,
      float* out,
      const std::int32_t* compressed_indices_table,
      const int* mask);
};

template <
    typename inType,
    typename indxType,
    typename offsetType,
    inst_set_t instSet,
    bool ROWWISE_SPARSE = false>
class GenEmbeddingSpMDMLookup {
 public:
  GenEmbeddingSpMDMLookup() {}
  typename ReturnFunctionSignature<
      inType,
      indxType,
      offsetType,
      ROWWISE_SPARSE>::jit_embedding_kernel
  getOrCreate(
      int block_size,
      bool has_weight,
      bool is_weight_positional,
      bool normalize_by_lengths,
      int prefetch,
      bool use_offsets);

 private:
  static asmjit::JitRuntime& runtime() {
    static asmjit::JitRuntime rt; //< JIT Runtime for asmjit,
                                  // depents on other static
                                  // variables.  Required to prevent
                                  // initialization order fiasco
    return rt;
  }

  static std::mutex rtMutex_; ///< Controll access to runtime;

  // The hash depends on embedding dimension (block size), weighted sls,
  // positional weights, normalize by lenths, prefetch distance, and use_offsets
  static CodeCache<
      std::tuple<int, bool, bool, bool, int, bool>,
      typename ReturnFunctionSignature<
          inType,
          indxType,
          offsetType,
          ROWWISE_SPARSE>::jit_embedding_kernel>
      codeCache_; ///< JIT Code Cache for reuse.
}; // GenEmbeddingSpmDMLookup

template <
    typename inType,
    typename indxType,
    typename offsetType,
    inst_set_t instSet,
    bool ROWWISE_SPARSE>
std::mutex GenEmbeddingSpMDMLookup<
    inType,
    indxType,
    offsetType,
    instSet,
    ROWWISE_SPARSE>::rtMutex_;

template <
    typename inType,
    typename indxType,
    typename offsetType,
    inst_set_t instSet,
    bool ROWWISE_SPARSE>
CodeCache<
    std::tuple<int, bool, bool, bool, int, bool>,
    typename ReturnFunctionSignature<
        inType,
        indxType,
        offsetType,
        ROWWISE_SPARSE>::jit_embedding_kernel>
    GenEmbeddingSpMDMLookup<
        inType,
        indxType,
        offsetType,
        instSet,
        ROWWISE_SPARSE>::codeCache_;

template <
    typename inType,
    typename indxType,
    typename offsetType,
    inst_set_t instSet,
    bool ROWWISE_SPARSE>
typename ReturnFunctionSignature<
    inType,
    indxType,
    offsetType,
    ROWWISE_SPARSE>::jit_embedding_kernel
GenEmbeddingSpMDMLookup<inType, indxType, offsetType, instSet, ROWWISE_SPARSE>::
    getOrCreate(
        int block_size,
        bool has_weight,
        bool is_weight_positional,
        bool normalize_by_lengths,
        int prefetch,
        bool use_offsets) {
  std::tuple<int, bool, bool, bool, int, bool> kernelSig = std::make_tuple(
      block_size,
      has_weight,
      is_weight_positional,
      normalize_by_lengths,
      prefetch,
      use_offsets);

  return codeCache_.getOrCreate(
      kernelSig,
      [&]() -> typename ReturnFunctionSignature<
                inType,
                indxType,
                offsetType,
                ROWWISE_SPARSE>::jit_embedding_kernel {
        bool is8bit = std::is_same<inType, std::uint8_t>::value;
        bool is16bit = std::is_same<inType, float16>::value;

        // TODO: Make this tunable
        int pref_dist = prefetch;
        bool areIndices64b = std::is_same<indxType, std::int64_t>::value;

        asmjit::CodeHolder code;
        code.init(runtime().codeInfo());
        x86::Assembler assembler(&code);
        x86::Emitter* a = assembler.as<x86::Emitter>();
#if defined(FBGEMM_LOG_CODE)
        std::string filename = "embeddinglookup";
        if (is8bit) {
          filename += "_8bit";
        } else if (is16bit) {
          filename += "_fp16";
        }
        filename += "_emd_dim_" + std::to_string(block_size);
        filename += areIndices64b ? "_64bit" : "_32bit";
        filename += instSet == inst_set_t::avx512 ? "_avx512" : "_avx2";
        if (prefetch) {
          filename += "_prefetch";
        }
        if (has_weight) {
          filename += "_hasweight";
        }
        if (normalize_by_lengths) {
          filename += "_normalize_by_lengths";
        }
        if (!use_offsets) {
          filename += "_use_lengths";
        }
        if (ROWWISE_SPARSE) {
          filename += "_rowwise_sparse";
        }
        filename += ".txt";
        FILE* codeLogFile = fopen(filename.c_str(), "w");
        asmjit::FileLogger* codeLogger = new asmjit::FileLogger(codeLogFile);
        code.setLogger(codeLogger);
#endif
        // arguments to the function created
        x86::Gp output_size = a->zdi();
        // index_size will be overwritten to hold the end address of indices
        x86::Gp index_size = a->zsi();
        x86::Gp data_size = a->zdx();
        x86::Gp input = a->zcx();
        int reg_id = 8;
        x86::Gp indices = a->gpz(reg_id); // 8
        ++reg_id;
        x86::Gp lengths = a->gpz(reg_id); // 9
        ++reg_id;
        x86::Gp weights = a->gpz(reg_id); // 10
        ++reg_id;
        x86::Gp out = a->gpz(reg_id); // 11

        x86::Gp compressed_indices_table;
        if (ROWWISE_SPARSE) {
          ++reg_id;
          compressed_indices_table = a->gpz(reg_id); // 12
        }
        ++reg_id;
        x86::Gp scratchReg1_ = a->gpz(reg_id); // 12 or 13, also for mask

        ++reg_id;
        x86::Gpd lengths_R_ = a->gpz(reg_id).r32(); // 13 or 14
        ++reg_id;
        x86::Gp scratchReg2_ = a->gpz(reg_id); // 14 or 15

        asmjit::FuncDetail func;

        if (ROWWISE_SPARSE) {
          func.init(
              asmjit::FuncSignatureT<
                  bool,
                  std::int64_t, // output_size
                  std::int64_t, // index_size
                  std::int64_t, // uncompressed_data_size
                  const inType*, // input uint8_t or float
                  const indxType*, // indices
                  const offsetType*, // offsets or lengths
                  const float*, // weights
                  float*, // out
                  const std::int32_t*, // compressed_indices_table and then mask
                  const int*>(asmjit::CallConv::kIdHost));
        } else {
          func.init(asmjit::FuncSignatureT<
                    bool,
                    std::int64_t, // output_size
                    std::int64_t, // index_size
                    std::int64_t, // data_size
                    const inType*, // input uint8_t or float
                    const indxType*, // indices
                    const offsetType*, // offsets or lengths
                    const float*, // weights
                    float*, // out and then mask
                    const int*>(asmjit::CallConv::kIdHost));
        }

        asmjit::FuncFrame frame;
        frame.init(func);

        if (instSet == inst_set_t::avx2) {
          frame.setDirtyRegs(
              x86::Reg::kGroupVec,
              asmjit::Support::bitMask(0, 1, 2, 3, 4, 5, 6, 7) |
                  asmjit::Support::bitMask(8, 9, 10, 11, 12, 13, 14, 15));
        } else {
          frame.setDirtyRegs(
              x86::Reg::kGroupVec,
              asmjit::Support::bitMask(0, 1, 2, 3, 4, 5, 6, 7) |
                  asmjit::Support::bitMask(8, 9, 10, 11, 12, 13, 14, 15) |
                  asmjit::Support::bitMask(16, 17, 18, 19, 20, 21, 22, 23) |
                  asmjit::Support::bitMask(24, 25, 26, 27, 28, 29, 30, 31));
        }

        frame.setDirtyRegs(
            x86::Reg::kGroupGp,
            reg_id == 15
                ? asmjit::Support::bitMask(8, 9, 10, 11, 12, 13, 14, 15)
                : asmjit::Support::bitMask(8, 9, 10, 11, 12, 13, 14));

        asmjit::FuncArgsAssignment args(&func);
        if (ROWWISE_SPARSE) {
          args.assignAll(
              output_size,
              index_size,
              data_size,
              input,
              indices,
              lengths,
              weights,
              out,
              compressed_indices_table,
              scratchReg1_);
        } else {
          args.assignAll(
              output_size,
              index_size,
              data_size,
              input,
              indices,
              lengths,
              weights,
              out,
              scratchReg1_);
        }

        args.updateFuncFrame(frame);
        frame.finalize();

        a->emitProlog(frame);
        a->emitArgsAssignment(frame, args);

        constexpr int vlen = simd_info<instSet>::WIDTH_32BIT_ELEMS;
        constexpr int NUM_VEC_REG = simd_info<instSet>::NUM_VEC_REGS;
        int unroll_factor = NUM_VEC_REG;

        typedef typename simd_info<instSet>::vec_reg_t vec_reg_t;

        int num_vec_regs_per_block = (block_size + vlen - 1) / vlen;
        int remainder = block_size % vlen;

        vec_reg_t scale_vreg; // holds scale
        vec_reg_t bias_vreg; // holds bias
        vec_reg_t w_vreg; // for weighted sls -- weights
        vec_reg_t
            vlen_inv_vreg; // used for normalize by lengths -- 1/ lengths[i]
        vec_reg_t src_vreg; // for holding embedding value temporarily
        x86::Ymm mask_vreg; // mask for avx2
        x86::Xmm mask_fp16_vreg; // mask for loading fp16 in avx2

        if (is8bit) {
          // We need 2 vec registers for 1. scale 2. bias
          --unroll_factor;
          scale_vreg = vec_reg_t(unroll_factor);
          --unroll_factor;
          bias_vreg = vec_reg_t(unroll_factor);
        }

        if (is8bit || is16bit || (remainder && instSet == inst_set_t::avx2)) {
          --unroll_factor;
          src_vreg = vec_reg_t(unroll_factor);
        }

        if (has_weight) {
          --unroll_factor;
          w_vreg = vec_reg_t(unroll_factor);
        }

        if (remainder && instSet == inst_set_t::avx2) {
          // AVX512 doesn't need to use vector register for masking
          --unroll_factor;
          mask_vreg = x86::ymm(unroll_factor);
          if (remainder > 1 && is16bit) {
            --unroll_factor;
            mask_fp16_vreg = x86::xmm(unroll_factor);
          }
        }

        if (normalize_by_lengths) {
          --unroll_factor;
          vlen_inv_vreg = vec_reg_t(unroll_factor);
        }

        if (remainder) {
          if (instSet == inst_set_t::avx2) {
            a->vmovups(
                mask_vreg,
                x86::ymmword_ptr(
                    scratchReg1_, (vlen - remainder) % vlen * sizeof(int32_t)));
            if (remainder > 1 && is16bit) {
              a->vmovups(
                  mask_fp16_vreg,
                  x86::xmmword_ptr(
                      scratchReg1_,
                      (vlen / 2 - remainder / 2) % (vlen / 2) *
                          sizeof(int32_t)));
              // We need to keep using the stack during the main loop
              a->lea(
                  x86::rsp,
                  x86::dword_ptr(
                      x86::rsp, static_cast<int32_t>(-vlen * sizeof(int32_t))));
            }
          } else {
            a->mov(scratchReg1_, (1 << remainder) - 1);
            a->kmovw(x86::k(1), scratchReg1_);
          }
        }

        // Compute the end address of indices
        a->imul(
            scratchReg1_,
            index_size,
            static_cast<asmjit::Imm>(sizeof(indxType)));
        a->add(scratchReg1_, indices);
        a->mov(index_size, scratchReg1_);

        asmjit::Label exit = a->newLabel();
        asmjit::Label error = a->newLabel();
        asmjit::Label LoopRangeIndexBegin = a->newLabel();
        asmjit::Label LoopRangeIndexEnd = a->newLabel();

        // rangeIndex loop begins (iterate output_size times)
        a->bind(LoopRangeIndexBegin);
        a->dec(output_size);
        a->jl(LoopRangeIndexEnd);

        if (normalize_by_lengths) {
          asmjit::Label IfLengthsBegin = a->newLabel();
          asmjit::Label IfLengthsEnd = a->newLabel();
          a->bind(IfLengthsBegin);
          if (use_offsets) {
            a->mov(lengths_R_, x86::dword_ptr(lengths, sizeof(offsetType)));
            a->sub(lengths_R_, x86::dword_ptr(lengths));
          } else {
            a->mov(lengths_R_, x86::dword_ptr(lengths));
          }
          a->cmp(lengths_R_, 1);
          // Initialize vlen_inv as 0 in case lengths is 0
          a->vxorps(vlen_inv_vreg, vlen_inv_vreg, vlen_inv_vreg);
          a->jl(IfLengthsEnd);

          if (instSet == inst_set_t::avx2) {
            x86::Xmm vlen_inv_vreg_xmm(vlen_inv_vreg.id());

            a->mov(scratchReg1_, 1);
            a->cvtsi2ss(vlen_inv_vreg_xmm, scratchReg1_);
            a->cvtsi2ss(x86::xmm0, lengths_R_);
            a->divss(vlen_inv_vreg_xmm, x86::xmm0);
            a->vpbroadcastd(vlen_inv_vreg, vlen_inv_vreg_xmm);
          } else { // avx512
            vec_reg_t temp_zmm = vec_reg_t(0);
            a->mov(scratchReg1_, 1);
            a->cvtsi2ss(x86::xmm(temp_zmm.id()), scratchReg1_);
            a->vpbroadcastd(vlen_inv_vreg, x86::xmm(temp_zmm.id()));
            a->vpbroadcastd(temp_zmm, lengths_R_);
            a->vcvtdq2ps(temp_zmm, temp_zmm);
            a->vdivps(vlen_inv_vreg, vlen_inv_vreg, temp_zmm);
          }
          a->bind(IfLengthsEnd);
        }

        for (int vec_idx = 0; vec_idx < num_vec_regs_per_block;
             vec_idx += unroll_factor) {
          int cur_unroll_factor =
              std::min(unroll_factor, num_vec_regs_per_block - vec_idx);

          // Initialize output regs
          for (int v = 0; v < cur_unroll_factor; ++v) {
            vec_reg_t out_vreg = vec_reg_t(v);
            a->vxorps(out_vreg, out_vreg, out_vreg);
          }

          if (use_offsets) {
            a->mov(lengths_R_, x86::dword_ptr(lengths, sizeof(offsetType)));
            a->sub(lengths_R_, x86::dword_ptr(lengths));
          } else {
            a->mov(lengths_R_, x86::dword_ptr(lengths));
          }

          // Array out of bound check
          a->imul(
              scratchReg1_,
              lengths_R_,
              static_cast<asmjit::Imm>(sizeof(indxType)));

          a->add(scratchReg1_, indices);
          a->cmp(scratchReg1_, index_size);
          a->jg(error);

          asmjit::Label LoopDataIndexBegin = a->newLabel();
          asmjit::Label LoopDataIndexEnd = a->newLabel();

          // dataIndex loop begins (iterate lengths_R_ times)
          a->bind(LoopDataIndexBegin);
          a->dec(lengths_R_);
          a->jl(LoopDataIndexEnd);

          // Array out of bound check
          if (areIndices64b) {
            a->mov(scratchReg1_, x86::qword_ptr(indices));
          } else {
            a->mov(scratchReg1_.r32(), x86::dword_ptr(indices));
          }
          // A trick to check x >= data_size or x < 0 in one shot by treating
          // scratchReg1_ as if it has unsigned value
          // (https://stackoverflow.com/a/34072155).
          a->cmp(scratchReg1_, data_size);
          a->jae(error);

          if (ROWWISE_SPARSE) {
            a->mov(
                scratchReg1_.r32(),
                x86::dword_ptr(
                    compressed_indices_table,
                    scratchReg1_,
                    2)); // use of 2 is to multiply by 4
          }

          int fused_block_size = is8bit
              ? block_size * sizeof(uint8_t) + 2 * sizeof(float)
              : block_size * sizeof(inType);

          if (pref_dist) {
            asmjit::Label pref_dist_reset_start = a->newLabel();
            asmjit::Label pref_dist_reset_end = a->newLabel();
            // out of bound handling for prefetch
            a->mov(scratchReg2_, indices);
            a->add(
                scratchReg2_,
                static_cast<asmjit::Imm>(pref_dist * sizeof(indxType)));
            a->cmp(scratchReg2_, index_size);
            a->jge(pref_dist_reset_start);

            if (areIndices64b) {
              a->mov(
                  scratchReg2_,
                  x86::qword_ptr(indices, pref_dist * sizeof(indxType)));
            } else {
              a->mov(
                  scratchReg2_.r32(),
                  x86::dword_ptr(indices, pref_dist * sizeof(indxType)));
            }

            a->jmp(pref_dist_reset_end);

            a->bind(pref_dist_reset_start);
            // things are not okay just get the current row
            // this can be improved to getting the max dist row.
            if (areIndices64b) {
              a->mov(scratchReg2_, x86::qword_ptr(indices));
            } else {
              a->mov(scratchReg2_.r32(), x86::dword_ptr(indices));
            }

            a->bind(pref_dist_reset_end);
            if (ROWWISE_SPARSE) {
              asmjit::Label rowwise_sparse_pref_corner_case_begin =
                  a->newLabel();
              asmjit::Label rowwise_sparse_pref_corner_case_end = a->newLabel();
              a->cmp(scratchReg2_, data_size);
              a->jae(rowwise_sparse_pref_corner_case_begin);

              a->mov(
                  scratchReg2_.r32(),
                  x86::dword_ptr(
                      compressed_indices_table,
                      scratchReg2_,
                      2)); // use of 2 is to multiply by 4
              a->test(scratchReg2_.r32(), scratchReg2_.r32());
              // Check negative
              a->jns(rowwise_sparse_pref_corner_case_end);

              a->bind(rowwise_sparse_pref_corner_case_begin);
              // For corner case, just set prefetch row id to 0.
              a->xor_(scratchReg2_.r32(), scratchReg2_.r32());
              a->bind(rowwise_sparse_pref_corner_case_end);
            }
            a->imul(scratchReg2_, static_cast<asmjit::Imm>(fused_block_size));
          }

          a->add(indices, static_cast<asmjit::Imm>(sizeof(indxType)));

          if (has_weight) {
            a->vbroadcastss(w_vreg, x86::dword_ptr(weights));
            a->add(weights, static_cast<asmjit::Imm>(sizeof(float)));
          }

          if (ROWWISE_SPARSE) {
            a->cmp(scratchReg1_.r32(), static_cast<asmjit::Imm>(-1));
            a->je(LoopDataIndexBegin);
          }

          a->imul(scratchReg1_, static_cast<asmjit::Imm>(fused_block_size));

          // broadcast the scale
          x86::Mem scale_src, bias_src;
          constexpr int CACHE_LINE_LEN = 64;
          if (is8bit) {
            scale_src = x86::dword_ptr(
                input, scratchReg1_, 0, block_size * sizeof(uint8_t));
            bias_src = x86::dword_ptr(
                input,
                scratchReg1_,
                0,
                block_size * sizeof(uint8_t) + sizeof(float));
            a->vbroadcastss(scale_vreg, scale_src);
            a->vbroadcastss(bias_vreg, bias_src);

            if (pref_dist && fused_block_size % CACHE_LINE_LEN > 0 &&
                fused_block_size % CACHE_LINE_LEN <= 2 * sizeof(float)) {
              a->prefetcht0(x86::dword_ptr(
                  input,
                  scratchReg2_,
                  0,
                  fused_block_size / CACHE_LINE_LEN * CACHE_LINE_LEN));
            }
          }

          if (has_weight && is8bit) {
            a->vmulps(scale_vreg, scale_vreg, w_vreg);
            a->vmulps(bias_vreg, bias_vreg, w_vreg);
          }

          // The main computation
          for (int v = 0; v < cur_unroll_factor; ++v) {
            constexpr int BYTES_PER_VLOAD = vlen * sizeof(inType);
            auto src_addr = x86::dword_ptr(
                input, scratchReg1_, 0, (vec_idx + v) * BYTES_PER_VLOAD);
            vec_reg_t out_vreg = vec_reg_t(v);

            // For 8bit SLS convert usigned 8-bit to 32bit int, then to float
            // multiply with scale and then add with bias
            if (is8bit) {
              if (remainder && vec_idx + v == num_vec_regs_per_block - 1 &&
                  instSet == inst_set_t::avx512) {
                a->k(x86::k(1)).z().vpmovzxbd(src_vreg, src_addr);
              } else {
                // We don't use a mask for AVX2 since we can use the extra
                // "padding" of the 2 floats (= 8 chars) scale and bias
                // this ensures we never access out of bound data
                a->vpmovzxbd(src_vreg, src_addr);
              }
              a->vcvtdq2ps(src_vreg, src_vreg);
              a->vaddps(out_vreg, out_vreg, bias_vreg);
              a->vfmadd231ps(out_vreg, src_vreg, scale_vreg);
            } else if (is16bit) {
              if (remainder && vec_idx + v == num_vec_regs_per_block - 1) {
                if (instSet == inst_set_t::avx2) {
                  if (remainder % 2 == 0) {
                    a->vmaskmovps(
                        x86::xmm(src_vreg.id()), mask_fp16_vreg, src_addr);
                  } else {
                    a->vpbroadcastw(
                        x86::xmm(src_vreg.id()),
                        x86::word_ptr(
                            input,
                            scratchReg1_,
                            0,
                            (vec_idx + v) * BYTES_PER_VLOAD +
                                (remainder - 1) * sizeof(inType)));
                    if (remainder > 1) {
                      // AVX2 can't do masking for the last 16-bit so we store
                      // them to a stack and reload.
                      // First put broadcasted last 16-bit element
                      a->vmovups(
                          x86::xmmword_ptr(x86::rsp), x86::xmm(src_vreg.id()));
                      // Mask store the remaining 16-bit elements
                      a->vmaskmovps(
                          x86::xmm(src_vreg.id()), mask_fp16_vreg, src_addr);
                      a->vmaskmovps(
                          x86::xmmword_ptr(x86::rsp),
                          mask_fp16_vreg,
                          x86::xmm(src_vreg.id()));
                      // Load combined 16-bit elements
                      a->vmovups(
                          x86::xmm(src_vreg.id()), x86::xmmword_ptr(x86::rsp));
                    } // remainder > 1
                  } // remainder % 2
                  a->vcvtph2ps(
                      x86::ymm(src_vreg.id()), x86::xmm(src_vreg.id()));
                } else {
                  // avx512
                  a->k(x86::k(1)).z().vcvtph2ps(src_vreg, src_addr);
                }
              } else {
                // no remainder
                a->vcvtph2ps(src_vreg, src_addr);
              }
              if (has_weight) {
                a->vfmadd231ps(out_vreg, w_vreg, src_vreg);
              } else {
                a->vaddps(out_vreg, out_vreg, src_vreg);
              }
            } else {
              // This part for FP32 SLS
              if (remainder && vec_idx + v == num_vec_regs_per_block - 1 &&
                  instSet == inst_set_t::avx2) {
                a->vmaskmovps(
                    x86::ymm(src_vreg.id()),
                    x86::ymm(mask_vreg.id()),
                    src_addr);
              }
              if (has_weight) {
                if (remainder && vec_idx + v == num_vec_regs_per_block - 1) {
                  if (instSet == inst_set_t::avx2) {
                    a->vfmadd231ps(out_vreg, w_vreg, src_vreg);
                  } else {
                    a->k(x86::k(1)).vfmadd231ps(out_vreg, w_vreg, src_addr);
                  }
                } else {
                  a->vfmadd231ps(out_vreg, w_vreg, src_addr);
                }
              } else {
                if (remainder && vec_idx + v == num_vec_regs_per_block - 1) {
                  if (instSet == inst_set_t::avx2) {
                    a->vaddps(out_vreg, out_vreg, src_vreg);
                  } else {
                    a->k(x86::k(1)).vaddps(out_vreg, out_vreg, src_addr);
                  }
                } else {
                  a->vaddps(out_vreg, out_vreg, src_addr);
                }
              }
            }

            constexpr int VLOAD_PER_CACHE_LINE =
                CACHE_LINE_LEN / BYTES_PER_VLOAD;
            if (pref_dist && (vec_idx + v) % VLOAD_PER_CACHE_LINE == 0) {
              a->prefetcht0(x86::dword_ptr(
                  input, scratchReg2_, 0, (vec_idx + v) * BYTES_PER_VLOAD));
            }
          }

          a->jmp(LoopDataIndexBegin);
          a->bind(LoopDataIndexEnd);

          // This loop is for writing back out_vreg (results)
          // back to memory
          for (int v = 0; v < cur_unroll_factor; ++v) {
            auto dst_addr =
                x86::dword_ptr(out, (vec_idx + v) * vlen * sizeof(float));
            vec_reg_t out_vreg = vec_reg_t(v);

            if (normalize_by_lengths) {
              a->vmulps(out_vreg, out_vreg, vlen_inv_vreg);
            }

            if (remainder && vec_idx + v == num_vec_regs_per_block - 1) {
              if (instSet == inst_set_t::avx2) {
                a->vmaskmovps(dst_addr, mask_vreg, x86::Ymm(out_vreg.id()));
              } else {
                a->k(x86::k(1)).vmovups(dst_addr, out_vreg);
              }
            } else {
              a->vmovups(dst_addr, out_vreg);
            }
          }

          if (vec_idx + unroll_factor < num_vec_regs_per_block ||
              (has_weight && is_weight_positional)) {
            // Reset lengths_R_, indices, weights to run the dataIndex loop
            // again
            if (use_offsets) {
              a->mov(lengths_R_, x86::dword_ptr(lengths, sizeof(offsetType)));
              a->sub(lengths_R_, x86::dword_ptr(lengths));
            } else {
              a->mov(lengths_R_, x86::dword_ptr(lengths));
            }

            if (has_weight) {
              a->imul(
                  scratchReg1_,
                  lengths_R_,
                  static_cast<asmjit::Imm>(sizeof(float)));
              a->sub(weights, scratchReg1_);

              if (vec_idx + unroll_factor < num_vec_regs_per_block) {
                a->imul(
                    scratchReg1_,
                    static_cast<asmjit::Imm>(sizeof(indxType) / sizeof(float)));
                a->sub(indices, scratchReg1_);
              }
            } else {
              a->imul(
                  scratchReg1_,
                  lengths_R_,
                  static_cast<asmjit::Imm>(sizeof(indxType)));
              a->sub(indices, scratchReg1_);
            }
          }
        }

        a->add(lengths, static_cast<asmjit::Imm>(sizeof(offsetType)));
        a->add(out, static_cast<asmjit::Imm>(block_size * sizeof(float)));

        a->jmp(LoopRangeIndexBegin);
        a->bind(LoopRangeIndexEnd);

        a->cmp(indices, index_size);
        a->jne(error);
        a->mov(x86::eax, true);
        a->jmp(exit);
        a->bind(error);
        a->mov(x86::eax, false);
        a->bind(exit);

        if (remainder > 1 && instSet == inst_set_t::avx2 && is16bit) {
          a->lea(x86::rsp, x86::ymmword_ptr(x86::rsp, vlen * sizeof(int32_t)));
        }

        a->emitEpilog(frame);

        // jit_fused8bitembedding_kernel fn;
        typename ReturnFunctionSignature<
            inType,
            indxType,
            offsetType,
            ROWWISE_SPARSE>::jit_embedding_kernel fn;
        asmjit::Error err;
        {
          std::unique_lock<std::mutex> lock(rtMutex_);
          err = runtime().add(&fn, &code);
        }
        if (err) {
          std::cout << "Error: in fn add" << std::endl;
          return nullptr;
        }

#if defined(FBGEMM_LOG_CODE)
        fclose(codeLogFile);
        delete codeLogger;
#endif
        return fn;
      });
}

} // namespace

template <typename inType, typename indxType, typename offsetType>
typename EmbeddingSpMDMKernelSignature<inType, indxType, offsetType>::Type
GenerateEmbeddingSpMDM(
    const std::int64_t block_size,
    bool has_weight,
    bool normalize_by_lengths,
    int prefetch,
    bool is_weight_positional,
    bool use_offsets) {
  if (!cpuinfo_initialize()) {
    throw std::runtime_error("Failed to initialize cpuinfo!");
  }
  if ((std::is_same<inType, float>::value ||
       std::is_same<inType, float16>::value) &&
      block_size == 1 && fbgemmHasAvx2Support()) {
    return
        [=](std::int64_t output_size,
            std::int64_t index_size,
            std::int64_t data_size,
            const inType* input,
            const indxType* indices,
            const offsetType* offsets_or_lengths,
            const float* weights, // optional, can be null for non-weighted sum
            float* out) {
          return internal::EmbeddingSpMDMBlockSize1_(
              output_size,
              index_size,
              data_size,
              input,
              indices,
              offsets_or_lengths,
              weights,
              normalize_by_lengths,
              out,
              is_weight_positional,
              use_offsets);
        };
  } else if (fbgemmHasAvx512Support()) {
    static GenEmbeddingSpMDMLookup<
        inType,
        indxType,
        offsetType,
        inst_set_t::avx512>
        kernel_generator;
    const auto original_func = kernel_generator.getOrCreate(
        block_size,
        has_weight,
        is_weight_positional,
        normalize_by_lengths,
        prefetch,
        use_offsets);
    return [=](std::int64_t output_size,
               std::int64_t index_size,
               std::int64_t data_size,
               const inType* input,
               const indxType* indices,
               const offsetType* offsets_or_lengths,
               const float* weights,
               float* out) {
      return original_func(
          output_size,
          index_size,
          data_size,
          input,
          indices,
          offsets_or_lengths,
          weights,
          out,
          nullptr /* mask not used in avx512 */);
    };
  } else if (fbgemmHasAvx2Support()) {
    static GenEmbeddingSpMDMLookup<
        inType,
        indxType,
        offsetType,
        inst_set_t::avx2>
        kernel_generator;
    const auto original_func = kernel_generator.getOrCreate(
        block_size,
        has_weight,
        is_weight_positional,
        normalize_by_lengths,
        prefetch,
        use_offsets);
    return [=](std::int64_t output_size,
               std::int64_t index_size,
               std::int64_t data_size,
               const inType* input,
               const indxType* indices,
               const offsetType* offsets_or_lengths,
               const float* weights,
               float* out) {
      return original_func(
          output_size,
          index_size,
          data_size,
          input,
          indices,
          offsets_or_lengths,
          weights,
          out,
          internal::avx2_ps_or_epi32_combined_mask);
    };
  } else {
#ifdef VLOG
    VLOG(0) << "AVX2 or AVX512 not found, taking the slow path";
#endif
    return [=](std::int64_t output_size,
               std::int64_t index_size,
               std::int64_t data_size,
               const inType* input,
               const indxType* indices,
               const offsetType* offsets_or_lengths,
               const float* weights,
               float* out) {
      return EmbeddingSpMDM_ref(
          block_size,
          output_size,
          index_size,
          data_size,
          input,
          indices,
          offsets_or_lengths,
          weights,
          normalize_by_lengths,
          out,
          is_weight_positional);
    };
  }
}

template <typename inType, typename indxType, typename offsetType>
typename EmbeddingSpMDMRowWiseSparseKernelSignature<
    inType,
    indxType,
    offsetType>::Type
GenerateEmbeddingSpMDMRowWiseSparse(
    const int64_t block_size,
    bool has_weight,
    bool normalize_by_lengths,
    int prefetch,
    bool is_weight_positional,
    bool use_offsets) {
  if (!cpuinfo_initialize()) {
    throw std::runtime_error("Failed to initialize cpuinfo!");
  }
  if (fbgemmHasAvx512Support()) {
    static GenEmbeddingSpMDMLookup<
        inType,
        indxType,
        offsetType,
        inst_set_t::avx512,
        /*rowwise_sparse=*/true>
        kernel_generator;
    const auto original_func = kernel_generator.getOrCreate(
        block_size,
        has_weight,
        is_weight_positional,
        normalize_by_lengths,
        prefetch,
        use_offsets);
    return [=](std::int64_t output_size,
               std::int64_t index_size,
               std::int64_t uncompressed_data_size,
               const inType* input,
               const indxType* indices,
               const offsetType* offsets_or_lengths,
               const float* weights,
               float* out,
               const std::int32_t* compressed_indices_table) {
      return original_func(
          output_size,
          index_size,
          uncompressed_data_size,
          input,
          indices,
          offsets_or_lengths,
          weights,
          out,
          compressed_indices_table,
          nullptr /* mask not used in avx512 */);
    };
  } else if (fbgemmHasAvx2Support()) {
    static GenEmbeddingSpMDMLookup<
        inType,
        indxType,
        offsetType,
        inst_set_t::avx2,
        /*rowwise_sparse=*/true>
        kernel_generator;
    const auto original_func = kernel_generator.getOrCreate(
        block_size,
        has_weight,
        is_weight_positional,
        normalize_by_lengths,
        prefetch,
        use_offsets);
    return [=](std::int64_t output_size,
               std::int64_t index_size,
               std::int64_t uncompressed_data_size,
               const inType* input,
               const indxType* indices,
               const offsetType* offsets_or_lengths,
               const float* weights,
               float* out,
               const std::int32_t* compressed_indices_table) {
      return original_func(
          output_size,
          index_size,
          uncompressed_data_size,
          input,
          indices,
          offsets_or_lengths,
          weights,
          out,
          compressed_indices_table,
          internal::avx2_ps_or_epi32_combined_mask);
    };
  } else {
#ifdef VLOG
    VLOG(0) << "AVX2 or AVX512 not found, taking the slow path";
#endif
    return
        [=](int64_t output_size,
            int64_t index_size,
            int64_t uncompressed_data_size,
            const inType* input,
            const indxType* indices,
            const offsetType* offsets_or_lengths,
            const float* weights, // optional, can be null for non-weighted sum
            float* out,
            const std::int32_t* compressed_indices_table) {
          return EmbeddingSpMDMRowWiseSparse_ref(
              block_size,
              output_size,
              index_size,
              uncompressed_data_size,
              // compressed_data_size,
              input,
              indices,
              compressed_indices_table,
              offsets_or_lengths,
              weights,
              normalize_by_lengths,
              out,
              is_weight_positional,
              use_offsets);
        };
  }
}

#define INSTANTIATE_SPMDM_BASE(NAME, IN_TYPE, INDEX_TYPE, OFFSET_TYPE)       \
  template FBGEMM_API                                                        \
      typename NAME##KernelSignature<IN_TYPE, INDEX_TYPE, OFFSET_TYPE>::Type \
          Generate##NAME<IN_TYPE, INDEX_TYPE, OFFSET_TYPE>(                  \
              const std::int64_t block_size,                                 \
              bool has_weight,                                               \
              bool normalize_by_lengths,                                     \
              int prefetch,                                                  \
              bool is_weight_positional,                                     \
              bool use_offsets);

#define INSTANTIATE_SPMDM_NAME(IN_TYPE, INDEX_TYPE, OFFSET_TYPE)           \
  INSTANTIATE_SPMDM_BASE(EmbeddingSpMDM, IN_TYPE, INDEX_TYPE, OFFSET_TYPE) \
  INSTANTIATE_SPMDM_BASE(                                                  \
      EmbeddingSpMDMRowWiseSparse, IN_TYPE, INDEX_TYPE, OFFSET_TYPE)

#define INSTANTIATE_SPMDM_OFFSET_T(IN_TYPE, INDEX_TYPE)     \
  INSTANTIATE_SPMDM_NAME(IN_TYPE, INDEX_TYPE, std::int32_t) \
  INSTANTIATE_SPMDM_NAME(IN_TYPE, INDEX_TYPE, std::int64_t)

#define INSTANTIATE_SPMDM_INDEX_T(IN_TYPE)          \
  INSTANTIATE_SPMDM_OFFSET_T(IN_TYPE, std::int32_t) \
  INSTANTIATE_SPMDM_OFFSET_T(IN_TYPE, std::int64_t)

INSTANTIATE_SPMDM_INDEX_T(float)
INSTANTIATE_SPMDM_INDEX_T(float16)
INSTANTIATE_SPMDM_INDEX_T(std::uint8_t)

#undef INSTANTIATE_SPMDM_INDEX_T
#undef INSTANTIATE_SPMDM_OFFSET_T
#undef INSTANTIATE_SPMDM_NAME
#undef INSTANTIATE_SPMDM_BASE

} // namespace fbgemm
